import os
import cv2
import numpy as np
from pdf2image import convert_from_path
from PIL import Image
import tempfile
import os.path
import random
import time
import argparse  # 用于命令行参数解析
from fastapi import APIRouter
import shutil  # 用于删除目录


def is_landscape(image):
    """判断图片是否为横向"""
    width, height = image.size
    return width > height

def split_image_vertically(image):
    """
    垂直拆分图片为两部分，智能检测空白区域进行切割
    切分位置满足特定条件：
    1. 切分位置严格限制在图像高度的40%-60%区间内
    2. 切分处上方一行应有较多空白
    3. 切分处下方一行不能只有个别位置有内容而大部分是空白
    """
    # 转换为OpenCV格式进行处理
    img_np = np.array(image)
    
    # 如果是RGB图像，转换为灰度图
    if len(img_np.shape) == 3:
        gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    else:
        gray = img_np
        
    height, width = gray.shape
    
    # 进行二值化处理，突出文本区域
    _, binary = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY_INV)
    
    # 严格限制搜索范围在图像高度的40%-60%区间内
    search_start = int(height * 0.4)
    search_end = int(height * 0.6)
    
    # 潜在切分点评分
    potential_split_points = []
    
    # 在搜索范围内逐行评估
    for row in range(search_start, search_end):
        # 检查当前行是否完全空白
        current_row_sum = np.sum(binary[row, :])
        if current_row_sum == 0:  # 完全空白行
            # 计算上方一行的空白程度（值越小表示空白越多）
            upper_row_sum = np.sum(binary[row-1, :]) if row > 0 else 0
            upper_row_nonzero = np.count_nonzero(binary[row-1, :]) if row > 0 else 0
            
            # 计算下方一行的内容分布情况
            if row + 1 < height:
                lower_row = binary[row+1, :]
                lower_row_sum = np.sum(lower_row)
                lower_row_nonzero = np.count_nonzero(lower_row)
                
                # 计算下方一行的内容聚集程度
                # 如果下方一行有内容，但内容分散（只有少数像素点有内容），则应避免
                if lower_row_nonzero > 0:
                    # 计算连续非零区域的数量（分段）
                    segments = 0
                    in_segment = False
                    for pixel in lower_row:
                        if pixel > 0 and not in_segment:
                            segments += 1
                            in_segment = True
                        elif pixel == 0:
                            in_segment = False
                    
                    # 计算下方一行的内容集中度
                    # 如果内容分散（段数多，但非零像素少），则不适合作为切分点
                    concentration = lower_row_nonzero / max(segments, 1)
                    
                    # 检查下方一行是否只有少数几处有内容
                    content_ratio = lower_row_nonzero / width
                    sparse_content = segments > 2 and content_ratio < 0.1
                    
                    # 如果下方一行内容不是太分散，且上方一行空白较多，则是个好的切分点
                    if not sparse_content and upper_row_sum < width * 0.3:
                        score = (width - upper_row_sum) * concentration  # 上方空白多且下方内容集中的得分高
                        potential_split_points.append((row, score))
            
            # 如果下方没有行了，则只考虑上方一行的空白程度
            elif upper_row_sum < width * 0.3:
                potential_split_points.append((row, width - upper_row_sum))
    
    # 如果找到了合适的切分点
    if potential_split_points:
        # 按得分降序排序
        potential_split_points.sort(key=lambda x: x[1], reverse=True)
        split_point = potential_split_points[0][0]
    else:
        # 如果在40%-60%区间内没有找到合适的空白行，则找到该区间内像素和最小的行
        row_sums = np.sum(binary[search_start:search_end, :], axis=1)
        min_row_index = np.argmin(row_sums)
        split_point = search_start + min_row_index
    
    # 确保拆分点绝对在40%-60%区间内
    if split_point < search_start:
        split_point = search_start
    elif split_point >= search_end:
        split_point = search_end - 1
    
    # 上半部分
    upper_part = image.crop((0, 0, width, split_point))
    
    # 下半部分
    lower_part = image.crop((0, split_point, width, height))
    
    return upper_part, lower_part

def extract_and_save_invoice_header(image, output_dir, prefix="header", header_ratio=0.25):
    """
    提取图片中的表头部分并保存(提取发票表头)
    表头区域为两条长横线下方约20-40像素处的单行内容，表头宽度根据内容区域确定
    
    参数:
        image: PIL图像对象
        output_dir: 输出目录
        prefix: 输出文件名前缀
        header_ratio: 表头占整个图像的高度比例（默认为25%，仅作为参考）
    
    返回:
        保存的表头图像路径
    """
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 获取图像尺寸
    width, height = image.size
    
    # 转换为OpenCV格式处理
    img_np = np.array(image)
    
    # 如果是RGB图像，转换为灰度图
    if len(img_np.shape) == 3:
        gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    else:
        gray = img_np
    
    # 进行二值化处理
    _, binary = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY_INV)
    
    # 水平投影，计算每行像素和
    row_sums = np.sum(binary, axis=1)
    
    # 使用形态学操作寻找水平线
    # 创建一个水平结构元素
    horizontal_size = width // 30
    horizontal_struct = cv2.getStructuringElement(cv2.MORPH_RECT, (horizontal_size, 1))
    
    # 应用形态学操作提取水平线
    horizontal = cv2.erode(binary, horizontal_struct)
    horizontal = cv2.dilate(horizontal, horizontal_struct)
    
    # 水平线所在行及其长度
    horizontal_lines = []
    for i in range(height):
        # 计算该行中的水平线长度
        row = horizontal[i, :]
        if np.sum(row) > 0:  # 如果行中有水平线像素
            # 找到连续的非零像素区域
            line_start = None
            line_end = None
            for j in range(width):
                if row[j] > 0 and line_start is None:
                    line_start = j
                elif row[j] == 0 and line_start is not None and line_end is None:
                    line_end = j - 1
                    break
            
            if line_start is not None:
                if line_end is None:  # 如果线一直延伸到图像边缘
                    line_end = width - 1
                
                line_length = line_end - line_start + 1
                line_ratio = line_length / width
                
                # 只记录较长的横线（占比超过50%）
                if line_ratio > 0.5:
                    horizontal_lines.append((i, line_start, line_end, line_ratio))
    
    # 按行号排序
    horizontal_lines.sort(key=lambda x: x[0])
    
    # 查找两条长横线
    long_lines = []
    prev_line = -100  # 初始化为一个不可能的行号
    for line_info in horizontal_lines:
        line, start, end, ratio = line_info
        # 如果与前一条线距离超过5像素，认为是不同的线
        if line - prev_line > 5:
            long_lines.append(line_info)
            prev_line = line
            
            # 如果已经找到两条长横线，停止查找
            if len(long_lines) == 2:
                break
    
    # 默认表头边界
    header_boundary = int(height * header_ratio)
    left_boundary = 0
    right_boundary = width
    
    # 如果找到了两条长横线
    if len(long_lines) >= 2:
        second_line = long_lines[1][0]  # 第二条横线的行号
        
        # 在第二条长横线下方20-40像素范围内寻找内容行
        search_start = second_line + 20
        search_end = min(second_line + 40, height)
        
        # 在指定范围内查找内容密集的行
        max_density = 0
        max_density_row = search_start
        content_left = width  # 初始化为最右边
        content_right = 0     # 初始化为最左边
        
        # 首先找到内容最密集的行
        for i in range(search_start, search_end):
            row_density = row_sums[i] / width
            if row_density > max_density:
                max_density = row_density
                max_density_row = i
        
        # 如果找到了内容行，确定内容的左右边界
        if max_density > 0.03:  # 有足够的内容
            # 分析行间距，找到单行的上下边界
            row_start = max_density_row
            row_end = max_density_row
            
            # 向上搜索，找到当前行的上边界（空白行）
            for i in range(max_density_row-1, search_start-1, -1):
                if row_sums[i] / width < 0.01:  # 如果是空白行
                    break
                row_start = i
            
            # 向下搜索，找到当前行的下边界（空白行或密度明显下降）
            for i in range(max_density_row+1, search_end+1):
                if i >= height or row_sums[i] / width < 0.01:  # 如果是空白行或超出图像范围
                    break
                # 检查是否是新的一行开始（密度明显下降后又上升）
                if (i > max_density_row+1 and 
                    row_sums[i] / width > row_sums[i-1] / width * 1.5):
                    break
                row_end = i
            
            # 确定单行内容的左右边界
            for i in range(row_start, row_end+1):
                row = binary[i, :]
                # 找到该行内容的左右边界
                nonzero_indices = np.nonzero(row)[0]
                if len(nonzero_indices) > 0:
                    row_left = nonzero_indices[0]
                    row_right = nonzero_indices[-1]
                    
                    # 更新整体内容的左右边界
                    content_left = min(content_left, row_left)
                    content_right = max(content_right, row_right)
            
            # 添加边距
            margin = int(width * 0.02)  # 2%的边距
            left_boundary = max(0, content_left - margin)
            right_boundary = min(width, content_right + margin)
            
            # 设置表头边界，确保只包含单行内容
            # 添加少量额外空间确保不截断
            header_boundary = row_end + 3  # 只包含当前行加少量额外空间
        else:
            # 如果没找到明显的内容行，使用默认位置和宽度
            header_boundary = second_line + 30
    
    # 提取表头部分，使用根据内容确定的左右边界
    header_image = image.crop((left_boundary, 0, right_boundary, header_boundary))
    
    # 生成随机文件名
    timestamp = int(time.time())
    random_num = random.randint(1000, 9999)
    filename = f"{prefix}_{timestamp}_{random_num}.png"
    file_path = os.path.join(output_dir, filename)
    
    # 保存表头图像
    header_image.save(file_path)
    
    print(f"表头已提取并保存到: {file_path}")
    return file_path

def add_header_to_image(image, header_image):
    """
    将表头添加到图像的顶部
    
    参数:
        image: 需要添加表头的PIL图像对象
        header_image: 表头的PIL图像对象
    
    返回:
        添加了表头的新图像
    """
    # 获取两个图像的尺寸
    img_width, img_height = image.size
    header_width, header_height = header_image.size
    
    # 调整表头宽度以匹配目标图像
    if header_width != img_width:
        ratio = img_width / header_width
        new_header_height = int(header_height * ratio)
        header_image = header_image.resize((img_width, new_header_height), Image.LANCZOS)
        header_height = new_header_height
    
    # 创建新图像，高度为两者之和
    new_img = Image.new('RGB', (img_width, img_height + header_height), (255, 255, 255))
    
    # 粘贴表头到顶部
    new_img.paste(header_image, (0, 0))
    
    # 粘贴原图像到表头下方
    new_img.paste(image, (0, header_height))
    
    return new_img

def splict_pdf_url(pdf_url, output_dir, add_header_to_second_page=False, save_header=False):
    """
    处理PDF中的图片，非横向图片则拆分
    
    参数:
        pdf_url: 输入PDF路径
        output_dir: 输出PDF目录
        add_header_to_second_page: 是否给截断后的第二页添加表头
        save_header: 是否保存提取的表头图像，默认为False
    """
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 获取输入文件名（不带路径和扩展名）
    input_filename = os.path.basename(pdf_url)
    input_filename_no_ext = os.path.splitext(input_filename)[0]
    
    # 生成随机数部分（使用时间戳和随机数结合）
    random_part = f"{int(time.time())}_{random.randint(1000, 9999)}"
    
    # 构建输出文件完整路径（添加随机数）
    output_pdf_path = os.path.join(output_dir, f"{input_filename_no_ext}_processed_{random_part}.pdf")
    
    # 创建用于保存表头的目录
    headers_dir = os.path.join(output_dir, "headers")
    if save_header and not os.path.exists(headers_dir):
        os.makedirs(headers_dir)
    
    # 在output_dir中创建临时目录，而不是使用系统临时目录
    temp_dir = os.path.join(output_dir, f"temp_{random_part}")
    if not os.path.exists(temp_dir):
        os.makedirs(temp_dir)
    
    try:
        # 将PDF转换为图片
        images = convert_from_path(pdf_url, output_folder=temp_dir, fmt='jpeg')
        
        processed_images = []
        header_paths = []
        header_images = []  # 存储每页的表头图像对象
        
        # 创建临时表头目录
        temp_header_dir = os.path.join(temp_dir, "temp_headers")
        if not os.path.exists(temp_header_dir):
            os.makedirs(temp_header_dir)
        
        for i, img in enumerate(images):
            # 如果需要保存表头或需要添加表头到第二页，则提取并保存表头
            if save_header or add_header_to_second_page:
                if save_header:
                    # 提取并保存表头到文件
                    header_path = extract_and_save_invoice_header(
                        img, 
                        headers_dir, 
                        prefix=f"header_page{i+1}_{input_filename_no_ext}"
                    )
                    header_paths.append(header_path)
                    header_img = Image.open(header_path)
                else:
                    # 提取表头但不保存到文件
                    header_path = extract_and_save_invoice_header(
                        img, 
                        temp_header_dir, 
                        prefix=f"temp_header_page{i+1}"
                    )
                    header_img = Image.open(header_path)
                    # 不添加到返回的路径列表中
                
                header_images.append(header_img)
            
            if is_landscape(img):
                # 横向图片，直接保留
                processed_images.append(img)
            else:
                # 纵向图片，拆分为两部分
                upper, lower = split_image_vertically(img)
                processed_images.append(upper)
                
                # 如果启用了添加表头选项，则给下半部分添加表头
                if add_header_to_second_page and header_images:
                    # 使用当前页面的表头
                    current_header_img = header_images[i]
                    lower_with_header = add_header_to_image(lower, current_header_img)
                    processed_images.append(lower_with_header)
                else:
                    processed_images.append(lower)
        
        # 保存处理后的图片为PDF
        if processed_images:
            processed_images[0].save(
                output_pdf_path, 
                save_all=True, 
                append_images=processed_images[1:]
            )
        
        print(f"处理完成，输出PDF已保存到: {output_pdf_path}")
        if save_header:
            print(f"提取的表头已保存到: {headers_dir}")
        
        # 关闭所有图像对象，确保没有文件被占用
        for img in processed_images:
            img.close()
        for img in header_images:
            img.close()
            
        return output_pdf_path, header_paths
    
    except Exception as e:
        print(f"处理PDF时发生错误: {str(e)}")
        raise e
    finally:
        # 无论处理成功还是失败，都清理临时目录
        try:
            if os.path.exists(temp_dir):
                shutil.rmtree(temp_dir)
                print(f"临时文件已删除: {temp_dir}")
        except Exception as e:
            print(f"删除临时文件时出错: {str(e)}")

def pdf_split_table_head(pdf_url, output_dir=None, add_header=True, save_header=False):
    """
    路由函数：处理PDF，将竖向页面拆分为两页，默认给第二页添加表头
    
    参数:
        pdf_url: 输入PDF路径
        output_dir: 输出目录，如果为None则使用默认目录
        add_header: 是否给截断后的第二页添加表头，默认为True
        save_header: 是否保存提取的表头图像，默认为False
        
    返回:
        processed_pdf_url: 处理后的PDF文件路径(本地路径)，处理失败时返回None
    """
    try:
        # 检查输入文件是否存在
        if not os.path.exists(pdf_url):
            print(f"错误：输入文件不存在 - {pdf_url}")
            return None
            
        # 如果未指定输出目录，则使用默认目录
        if output_dir is None:
            # 获取输入文件所在目录
            input_dir = os.path.dirname(pdf_url)
            if input_dir:
                output_dir = os.path.join(input_dir, "output_files")
            else:
                output_dir = "output_files"
        
        # 确保输出目录存在
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        # 处理PDF并获取结果
        processed_pdf_url, header_paths = splict_pdf_url(pdf_url, output_dir, add_header, save_header)
        
        # 返回处理后的PDF路径
        return processed_pdf_url
    except Exception as e:
        print(f"处理PDF时发生错误: {str(e)}")
        return None

# 创建FastAPI路由器
pdf_router = APIRouter(tags=["PDF处理"])

@pdf_router.post("/pdf-split/")
async def pdf_split_endpoint(pdf_url: str, add_header: bool = True, save_header: bool = False, output_dir: str = None):
    """
    处理PDF，将竖向页面拆分为两页
    
    参数:
    - pdf_url: PDF文件路径
    - add_header: 是否给截断后的第二页添加表头，默认为True
    - save_header: 是否保存提取的表头图像，默认为False
    - output_dir: 输出目录，如果为None则使用默认目录
    
    返回:
    - 处理后的PDF文件路径
    """
    result = pdf_split_table_head(pdf_url, output_dir, add_header, save_header)
    if result:
        return {"status": "success", "pdf_path": result}
    else:
        return {"status": "error", "message": "处理PDF失败"}

def test_pdf_split_table_head():
    # 获取输入输出路径和添加表头选项
    input_pdf = "input_files\(多页发票-小计-合计)1747278181013-dzfp_25952000000038488446_深圳市立创电子商务有限公司_20250331144352.pdf"
    
    result_pdf = pdf_split_table_head(input_pdf,output_dir="output_files",add_header=True,save_header=False)
    print("result_pdf:",result_pdf)

if __name__ == "__main__":
    test_pdf_split_table_head()